import os
import requests
from flask import Flask, render_template, request, jsonify
from config import Config
from werkzeug.utils import secure_filename

app = Flask(__name__)
app.config.from_object(Config)
Config.init_app(app)


def allowed_file(filename):
    return '.' in filename and filename.rsplit('.', 1)[1].lower() in app.config['ALLOWED_EXTENSIONS']


def generate_prompt(content):
    return f"""你是一个智能助手，请根据以下内容回答问题：
{content}
请用中文简洁明了地回答用户的问题。如果内容为空，请告知用户需要提供有效内容。"""


def query_ollama(prompt):
    try:
        response = requests.post(f"{app.config['OLLAMA_API_BASE']}/chat/completions",
                                 json={"model": app.config['MODEL_NAME'],
                                       "messages": [{"role": "user", "content": prompt}],
                                       "temperature": 0.7,
                                       # "stream": True,
                                       },
                                 timeout=300
                                 )
        response.raise_for_status()
        return response.json()['choices'][0]['message']['content']
    except Exception as e:
        return f"API请求失败: {str(e)}"


@app.route('/', methods=['GET'])
def index():
    return render_template('index.html')


@app.route('/process', methods=['POST'])
def process_input():
    text_input = request.form.get('text', '')
    file = request.files.get('file')
    file_content = ''

    if file and allowed_file(file.filename):
        filename = secure_filename(file.filename)
        filepath = os.path.join(app.config['UPLOAD_FOLDER'], filename)
        file.save(filepath)
        with open(filepath, 'r', encoding='utf-8') as f:
            file_content = f.read()
        os.remove(filepath)

    combined_input = text_input + "\n" + file_content
    if not combined_input.strip():
        return jsonify({'error': '请输入文本或上传文件'})

    prompt = generate_prompt(combined_input)
    result = query_ollama(prompt)
    return jsonify({'result': result})


if __name__ == '__main__':
    app.run(host='0.0.0.0', port=6789)
